package gwmongo

import (
	"context"
	"github.com/pkg/errors"
	"github.com/xiaonanln/goworld/engine/async"
	"github.com/xiaonanln/goworld/engine/gwlog"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"time"
)

const (
	_MONGODB_ASYNC_JOB_GROUP = "_mongodb"
)

var (
	errNoSession = errors.Errorf("no session, please dial")
)

// MongoDB is a MongoDB instance can be used to manipulate Mongo DBs
type DB struct {
	client *mongo.Client
	db     *mongo.Database
}

func (mdb *DB) checkConnected() bool {
	return mdb.client != nil && mdb.db != nil
}

// Dial connects the a MongoDB
// returns *MongoDB
func Dial(url string, dbname string, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		gwlog.Infof("Dialing MongoDB: %s ...", url)
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		client, err := mongo.Connect(ctx, options.Client().ApplyURI(url))
		if err != nil {
			return nil, err
		}
		db := client.Database(dbname)
		return &DB{client, db}, nil
	}, ac)
}

// Close closes MongoDB
func (mdb *DB) Close(ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		err := mdb.client.Disconnect(ctx)
		mdb.client = nil
		mdb.db = nil
		return nil, err
	}, ac)
}

// UseDB uses the specified DB
func (mdb *DB) UseDB(dbname string, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		mdb.db = mdb.client.Database(dbname)
		return nil, nil
	}, ac)
}

// SetMode sets the consistency mode
// Note: In the new driver, consistency mode is set differently. This is a placeholder.
func (mdb *DB) SetMode(consistency int, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		// In the new driver, consistency settings are more complex.
		// This is a placeholder and may need to be adjusted according to actual needs.
		return nil, nil
	}, ac)
}

// FindId finds document in collection by Id
func (mdb *DB) FindId(collectionName string, id interface{}, setupQuery func(query *options.FindOneOptions), ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		filter := bson.M{"_id": id}
		opts := options.FindOne()
		if setupQuery != nil {
			setupQuery(opts)
		}
		var res bson.M
		err := col.FindOne(context.TODO(), filter, opts).Decode(&res)
		return res, err
	}, ac)
}

// FindOne finds one document with specified query
func (mdb *DB) FindOne(collectionName string, query bson.M, setupQuery func(query *options.FindOneOptions), ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		opts := options.FindOne()
		if setupQuery != nil {
			setupQuery(opts)
		}
		var res bson.M
		err := col.FindOne(context.TODO(), query, opts).Decode(&res)
		return res, err
	}, ac)
}

// FindAll finds all documents with specified query
func (mdb *DB) FindAll(collectionName string, query bson.M, setupQuery func(query *options.FindOptions), ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		opts := options.Find()
		if setupQuery != nil {
			setupQuery(opts)
		}
		cursor, err := col.Find(context.TODO(), query, opts)
		if err != nil {
			return nil, err
		}
		var res []bson.M
		defer cursor.Close(context.TODO())
		for cursor.Next(context.TODO()) {
			var doc bson.M
			if err := cursor.Decode(&doc); err != nil {
				return nil, err
			}
			res = append(res, doc)
		}
		if err := cursor.Err(); err != nil {
			return nil, err
		}
		return res, nil
	}, ac)
}

// Count counts the number of documents by query
func (mdb *DB) Count(collectionName string, query bson.M, setupQuery func(query *options.CountOptions), ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		opts := options.Count()
		if setupQuery != nil {
			setupQuery(opts)
		}
		n, err := col.CountDocuments(context.TODO(), query, opts)
		return n, err
	}, ac)
}

// Insert inserts a document
func (mdb *DB) Insert(collectionName string, doc bson.M, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		_, err := col.InsertOne(context.TODO(), doc)
		return nil, err
	}, ac)
}

// InsertMany inserts multiple documents
func (mdb *DB) InsertMany(collectionName string, docs []bson.M, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		var insertDocs []interface{}
		for _, doc := range docs {
			insertDocs = append(insertDocs, doc)
		}
		_, err := col.InsertMany(context.TODO(), insertDocs)
		return nil, err
	}, ac)
}

// UpdateId updates a document by id
func (mdb *DB) UpdateId(collectionName string, id interface{}, update bson.M, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		filter := bson.M{"_id": id}
		_, err := col.UpdateOne(context.TODO(), filter, bson.M{"$set": update})
		return nil, err
	}, ac)
}

// Update updates a document by query
func (mdb *DB) Update(collectionName string, query bson.M, update bson.M, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		_, err := col.UpdateOne(context.TODO(), query, bson.M{"$set": update})
		return nil, err
	}, ac)
}

// UpdateAll updates all documents by query
func (mdb *DB) UpdateAll(collectionName string, query bson.M, update bson.M, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return 0, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		result, err := col.UpdateMany(context.TODO(), query, bson.M{"$set": update})
		if err != nil {
			return 0, err
		}
		return result.ModifiedCount, nil
	}, ac)
}

// UpsertId updates or inserts a document by id
func (mdb *DB) UpsertId(collectionName string, id interface{}, update bson.M, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		filter := bson.M{"_id": id}
		opts := options.Update().SetUpsert(true)
		result, err := col.UpdateOne(context.TODO(), filter, bson.M{"$set": update}, opts)
		if err != nil {
			return nil, err
		}
		return result.UpsertedID, nil
	}, ac)
}

// Upsert updates or inserts a document by query
func (mdb *DB) Upsert(collectionName string, query bson.M, update bson.M, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		opts := options.Update().SetUpsert(true)
		result, err := col.UpdateOne(context.TODO(), query, bson.M{"$set": update}, opts)
		if err != nil {
			return nil, err
		}
		return result.UpsertedID, nil
	}, ac)
}

// RemoveId removes a document by id
func (mdb *DB) RemoveId(collectionName string, id interface{}, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		filter := bson.M{"_id": id}
		_, err := col.DeleteOne(context.TODO(), filter)
		return nil, err
	}, ac)
}

// Remove removes a document by query
func (mdb *DB) Remove(collectionName string, query bson.M, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		_, err := col.DeleteOne(context.TODO(), query)
		return nil, err
	}, ac)
}

// RemoveAll removes all documents by query
func (mdb *DB) RemoveAll(collectionName string, query bson.M, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return 0, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		result, err := col.DeleteMany(context.TODO(), query)
		if err != nil {
			return 0, err
		}
		return result.DeletedCount, nil
	}, ac)
}

// EnsureIndex creates an index
func (mdb *DB) EnsureIndex(collectionName string, index mongo.IndexModel, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		_, err := col.Indexes().CreateOne(context.TODO(), index)
		return nil, err
	}, ac)
}

// EnsureIndexKey creates an index by keys
func (mdb *DB) EnsureIndexKey(collectionName string, keys []string, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		index := mongo.IndexModel{
			Keys: bson.D{},
		}
		for _, key := range keys {
			idxKeys, _ := index.Keys.(bson.D)
			idxKeys = append(idxKeys, bson.E{Key: key, Value: 1})
			index.Keys = idxKeys
		}
		_, err := col.Indexes().CreateOne(context.TODO(), index)
		return nil, err
	}, ac)
}

// DropIndex drops an index by keys
func (mdb *DB) DropIndex(collectionName string, keys []string, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		indexName := ""
		for _, key := range keys {
			if indexName != "" {
				indexName += "_"
			}
			indexName += key
		}
		_, err := col.Indexes().DropOne(context.TODO(), indexName)
		return nil, err
	}, ac)
}

// DropCollection drops a collection
func (mdb *DB) DropCollection(collectionName string, ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		col := mdb.db.Collection(collectionName)
		err := col.Drop(context.TODO())
		return nil, err
	}, ac)
}

// DropDatabase drops the database
func (mdb *DB) DropDatabase(ac async.AsyncCallback) {
	async.AppendAsyncJob(_MONGODB_ASYNC_JOB_GROUP, func() (interface{}, error) {
		if !mdb.checkConnected() {
			return nil, errNoSession
		}
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		err := mdb.db.Drop(ctx)
		return nil, err
	}, ac)
}
